import pandas as pd
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator


def DiscretizeByKmeans(k, data):
    """
    通过kmeans算法离散化
    :param k: 类别数
    :param data: 数据
    :return:
    """
    kmodel = KMeans(n_clusters=k, n_jobs=2)
    kmodel.fit(data.values.reshape(len(data), 1))
    print(kmodel.cluster_centers_)
    c = pd.DataFrame(kmodel.cluster_centers_).sort_values(0)
    w1 = c.rolling(2).mean().iloc[1:]
    w = [0] + list(w1[0]) + [data.max()]
    return pd.cut(data, w, labels=range(k))


def ClusterPlot(d, k, data):
    """
    画出离散化图
    :param d:
    :param k:
    :param data:
    :return:
    """
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号

    plt.figure(figsize=(8, 3))
    for j in range(0, k):
        plt.plot(data[d == j], [j for i in d[d == j]], 'o')

    plt.ylim(-1, k)
    ax = plt.gca()
    y_major_locator = MultipleLocator(1)
    ax.yaxis.set_major_locator(y_major_locator)
    return plt
